#!/usr/bin/python3

import argparse
import os
import subprocess
import difflib
import sys

RED = "\033[31m"
GREEN = "\033[32m"
YELLOW = "\033[33m"
ENDCOLOR = "\033[0m"

TESTDIR = "test"
tests = []

for item in os.scandir(TESTDIR):
    if item.is_dir():
        tests.append(os.path.basename(item.path))

CLASSMAP = {"func": "functional", "hfunc": "hidden_functional", "perf": "performance", "hperf": "hidden_performance"}

parser = argparse.ArgumentParser(description="测试脚本使用方法")
parser.add_argument(
    "--year", "-y", choices=tests + ["all"], default="2023", help="需要测试的年份"
)
parser.add_argument(
    "--types",
    "-t",
    choices=["func", "hfunc", "perf", "hperf"],
    default="func",
    nargs="+",
    help="需要测试的类别",
)

parser.add_argument("--single", "-s", nargs=1, help="测试单个文件")

parser.add_argument(
    "--mode",
    "-m",
    choices=["q", "soe", "detail"],
    default="q",
    help="q: 快速测试; soe: stop on error; detail: 发生错误时显示输出",
)
args = parser.parse_args()

cmake_ret = subprocess.run(
    ["cmake", "--build", "build", "--target", "sysyc", "--", "-j", "10"], check=True
)
sysyc = "build/bin/sysyc"
llvmlink = "llvm-link"
lli = "lli"
lib = "lib/sylib.ll"
tmpll = "tmp.ll"

test_years = args.year if isinstance(args.year, list) else [args.year]
test_types = args.types if isinstance(args.types, list) else [args.types]
test_single = args.single
test_mode = args.mode


def single(filepath):
    print(YELLOW + "COMPILING TEST" + ENDCOLOR + "\t" + filepath, end="")
    sys.stdout.flush()
    sysyc_res = subprocess.run(
        [sysyc, filepath], stdout=open(tmpll, "w"), text=True, stderr=subprocess.PIPE
    )
    sys.stdout.write("\33[2K\r")

    if sysyc_res.returncode != 0:
        print(RED + "COMPILING ERROR" + ENDCOLOR + "\t" + filepath)
        print(sysyc_res.stderr)
        return False

    print(YELLOW + "LINKING TEST" + ENDCOLOR + "\t" + filepath, end="")
    sys.stdout.flush()
    llvm_link_res = subprocess.run(
        [llvmlink, tmpll, lib, "-S", "-only-needed", "-o", tmpll],
        text=True,
        stderr=subprocess.PIPE,
    )
    sysystdin = ""
    if os.path.exists(filepath[:-2] + "in"):
        sysystdin = open(filepath[:-2] + "in", "r").read()
    sys.stdout.write("\33[2K\r")
    if llvm_link_res.returncode != 0:
        print(RED + "LINKING ERROR" + ENDCOLOR + "\t" + filepath)
        print(llvm_link_res.stderr)
        return False
    print(YELLOW + "RUNNING TEST" + ENDCOLOR + "\t" + filepath, end="")
    sys.stdout.flush()
    lli_process = subprocess.run(
        [lli, tmpll],
        input=sysystdin,
        stdout=subprocess.PIPE,
        text=True,
        stderr=subprocess.PIPE,
    )
    # 无需检测lli返回值, 因为main函数返回值可能不为0
    out = (lli_process.stdout.strip("\n") + "\n" + str(lli_process.returncode)).strip(
        "\n"
    )
    # 测试用例里out的构造疑似标准不一致, 主要问题体现在stdout的最后一个换行符是否需要清除, 因此特供两个版本, 总有一个对的
    another_out = (lli_process.stdout + str(lli_process.returncode)).strip(
        "\n"
    )
    expected_out = open(filepath[:-2] + "out", "r").read().strip("\n")
    sys.stdout.write("\33[2K\r")
    if out == expected_out or another_out == expected_out:
        print(GREEN + "ACCEPTED" + ENDCOLOR + "\t" + filepath)
        print(lli_process.stderr)
        return True
    else:
        print(RED + "RUNNING ERROR" + ENDCOLOR + "\t" + filepath)
        print(lli_process.stderr)
        if test_mode == "detail":
            # print(
            #     "".join(
            #         difflib.ndiff(
            #             expected_out.splitlines(keepends=True),
            #             out.splitlines(keepends=True),
            #         )
            #     )
            # )
            print(f"expected:\n{expected_out}\n\ngot:\n{out}\n\nanother:\n{another_out}\n")


total = 0
passed = 0
was = []

if test_single != None:
    single(test_single[0])
else:
    if test_years == ["all"]:
        test_years = tests
    test_types = [CLASSMAP[k] for k in test_types]
    for year in test_years:
        for type in test_types:
            dir = os.path.join(TESTDIR, year, type)
            for file in sorted(os.listdir(dir)):
                if file[-2:] != "sy":
                    continue
                total += 1
                if single(os.path.join(dir, file)):
                    passed += 1
                else:
                    was.append(os.path.join(dir, file))
    print(f"{passed}/{total} passed.")
    print(f"wrong list: {was}")

os.remove(tmpll)
